from datetime import datetime
from .position import Position
from data.providers.contract_provider import ContractProvider

class BacktestEngine:
    """回测引擎"""
    def __init__(self):
        self.data = None
        self.strategy = None
        self.position = Position()
        self.contract_provider = ContractProvider()
        self.current_contract = None
        self.next_contract = None
        self.switch_date = None
        
    def init_contract_data(self, product_code, start_date, end_date, months=[1,5,9], ahead_months=3):
        """初始化合约数据"""
        # 获取主力合约数据（用于生成信号）
        self.main_data = self.contract_provider.get_main_contract_data(
            product_code, start_date, end_date
        )
        
        # 获取实际交易合约数据
        contract_info = self.contract_provider.get_trading_contract_data(
            product_code,
            self.main_data['code'].iloc[0],  # 起始时刻的主力合约
            months,
            ahead_months,
            start_date,
            end_date
        )
        
        self.current_contract = contract_info['current_contract']
        self.next_contract = contract_info['next_contract']
        self.switch_date = contract_info['switch_date']
        self.contract_data = contract_info['contract_data']
        
    def on_contract_switch(self):
        """处理换月逻辑"""
        # 如果有持仓，等待平仓信号
        if self.position.has_position(self.current_contract):
            return
            
        # 没有持仓时，更新交易合约
        contract_info = self.contract_provider.get_trading_contract_data(
            self.product_code,
            self.current_contract,
            self.months,
            self.ahead_months
        )
        
        self.current_contract = contract_info['current_contract']
        self.next_contract = contract_info['next_contract']
        self.switch_date = contract_info['switch_date']
        
    def run(self):
        """运行回测"""
        for date, data in self.main_data.iterrows():
            # 生成信号
            signal = self.strategy.on_bar(data)
            
            # 检查是否需要换月
            if date >= self.switch_date:
                self.on_contract_switch()
            
            # 根据信号和当前交易合约处理交易
            if signal > 0:  # 做多信号
                if not self.position.has_position(self.current_contract):
                    self.position.open_long(self.current_contract, data['close'])
            elif signal < 0:  # 做空信号
                if self.position.has_position(self.current_contract):
                    self.position.close_all(self.current_contract, data['close']) 